'''
Author: xushaocong
Date: 2022-05-22 14:56:57
LastEditTime: 2022-05-27 12:28:40
LastEditors: xushaocong
Description: 
FilePath: /deep_motion_mag/data_loader.py
email: xushaocong@stu.xmu.edu.cn
'''


import os
import cv2
import numpy as np
import scipy.fftpack
import scipy.signal
from matplotlib import pyplot



import tensorflow as tf


def uint8_to_float(data):
  """
  converts numpy array of type "uint8"(default type when we load through opencv)
  to "float" datatype.
  """
  result = np.ndarray(shape=data.shape, dtype='float')
  result[:] = data * (1. / 255)
  return result

def get_capture_dimensions(capture):
    """Get the dimensions of a capture"""
    width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
    return width, height

def load_video(video_filename):
    """Load a video into a numpy array"""
    video_filename = str(video_filename)
    print("Loading " + video_filename)
    if not os.path.isfile(video_filename):
        raise Exception("File Not Found: %s" % video_filename)
    capture = cv2.VideoCapture(video_filename)
    frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    width, height = get_capture_dimensions(capture)
    fps = int(capture.get(cv2.CAP_PROP_FPS))
    x = 0
    vid_frames = np.zeros((frame_count, height, width, 3), dtype='uint8')
    while capture.isOpened():
        ret, frame = capture.read()
        if not ret:
            break

        vid_frames[x] = frame
        x += 1
    capture.release()

    return vid_frames, fps


def load_video_float(video_filename):
    """ Returns video as a numpy array of data type float"""
    vid_data, fps = load_video(video_filename)
    return uint8_to_float(vid_data), fps

'''
description: 
param {*} vid_data
param {*} fps
param {*} x_pixel
param {*} y_pixel
param {*} bounds
return {*}
'''
def show_frequencies(vid_data, fps, x_pixel=307, y_pixel=302, bounds=None):
    """Graph the average value of the video as well as the frequency strength"""
    averages = []
    size = vid_data.shape[0]
    if bounds:
        for x in range(1, vid_data.shape[0] - 1):
            averages.append(vid_data[x, bounds[2]:bounds[3], bounds[0]:bounds[1], :].sum())
    else:
        for x in range(1, vid_data.shape[0] - 1):
            averages.append(vid_data[x, x_pixel, y_pixel, :].sum())

    # averages = (averages-min(averages))/ (max(averages)-min(averages))
    averages = averages

    charts_x = 1
    charts_y = 2
    pyplot.figure(figsize=(20, 10))
    # pyplot.subplots_adjust(hspace=.7)

    pyplot.subplot(charts_y, charts_x, 1)
    pyplot.title(f"Pixel({x_pixel}, {y_pixel})")
    pyplot.xlabel("Time")
    pyplot.ylabel("Brightness")
    t1 = np.linspace(0,599,size-2) * 1/fps
    pyplot.plot(t1, averages)

    freqs = scipy.fftpack.fftfreq(len(averages), d=1.0 / fps)  # fps  #横坐标
    fft = abs(scipy.fftpack.fft(averages))
    idx = np.argsort(freqs)

    pyplot.subplot(charts_y, charts_x, 2)
    pyplot.title("FFT")
    pyplot.xlabel("Freq (Hz)")
    freqs = freqs[idx]
    fft = fft[idx]

    freqs = freqs[len(freqs) // 2 + 1:]
    fft = fft[len(fft) // 2 + 1:]
    y = abs(fft)
    pyplot.plot(freqs, y)
    max_indx = np.argmax(y)
    pyplot.plot(freqs[max_indx], y[max_indx], 'ks')
    show_max = '[' + str(freqs[max_indx]) + ' ' + str(y[max_indx]) + ']'
    pyplot.annotate(show_max, xytext=(freqs[max_indx], y[max_indx]), xy=(freqs[max_indx], y[max_indx]))

    # pyplot.show()
    pyplot.savefig("data/output/draw_frequence.jpg")
    
    
    
    
    
input_video_filename = 'data/output/data3_fl0.04_fh0.4_fs30.0_n2_differenceOfIIR/data3_fl0.04_fh0.4_fs30.0_n2_differenceOfIIR_259002.mp4'
vid, fps = load_video_float(input_video_filename)
show_frequencies(vid, 1000,  218,748)  #取点（307，302）